package com.example.ai.uploadImage;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.google.gson.JsonSyntaxException;
import com.jcraft.jsch.*;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Properties;

public class YOLOTrain {
    private static final String REMOTE_HOST  = "10.2.188.107";
    private static final int SSH_PORT  = 22;
    private static final String SSH_USER  = "root";
    private static final String SSH_PASSWORD  = "zjzhjs-2023";

    // 远程路径配置
    private static final String PYTHON_SCRIPT_PATH = "/root/yolo-project/py/yolo_detector.py";
    private static final String REMOTE_WORK_DIR = "/root/yolo-project/checkImage";

    public static void main(String[] args) {
        String localImagePath = "C:\\Users\\admin\\Desktop\\buildImage\\IMG20250000.jpg";
        JsonObject detectionResult = detectObjects(localImagePath);
        System.out.println("Detection Result: " + detectionResult);
    }

    public static JsonObject detectObjects(String localImagePath) {
        Session session = null;
        ChannelSftp sftpChannel = null;
        ChannelExec execChannel = null;

        try {
            // 1. 建立 SSH 连接
            JSch jsch = new JSch();
            session = jsch.getSession(SSH_USER, REMOTE_HOST, SSH_PORT);
            session.setPassword(SSH_PASSWORD);

            Properties config = new Properties();
            config.put("StrictHostKeyChecking", "no");
            session.setConfig(config);
            session.connect();

            // 2. 上传图像文件到服务器
            String remoteImageName = "upload_" + System.currentTimeMillis() + ".jpg";
            String remoteImagePath = REMOTE_WORK_DIR + "/" + remoteImageName;

            sftpChannel = (ChannelSftp) session.openChannel("sftp");
            sftpChannel.connect();
            sftpChannel.put(localImagePath, remoteImagePath);

            // 3. 执行 Python 检测脚本
            String command = String.format(
                    "cd /root/yolo-project && " +
                            "source /root/yolo11/bin/activate && " +
                    "python3 %s --image %s",
                    PYTHON_SCRIPT_PATH,
                    remoteImagePath
            );
            System.out.println(command);

            execChannel = (ChannelExec) session.openChannel("exec");
            execChannel.setCommand(command);

            // 同时捕获标准输出和错误输出
            ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
            ByteArrayOutputStream errorStream = new ByteArrayOutputStream();
            execChannel.setOutputStream(outputStream);
            execChannel.setErrStream(errorStream);

            execChannel.connect();

            // 等待命令执行完成
            while (!execChannel.isClosed()) {
                Thread.sleep(500);
            }

            // 检查退出状态
            int exitStatus = execChannel.getExitStatus();
            String output = outputStream.toString().trim();
            output = output.substring(output.indexOf("{"));
            String error = errorStream.toString().trim();

            if (exitStatus != 0) {
                throw new RuntimeException(
                        "Python脚本执行失败 (Exit " + exitStatus + ")\n" +
                                "错误信息: " + error + "\n" +
                                "输出内容: " + output
                );
            }

            // 4. 解析 JSON 结果（添加验证）
            if (output.isEmpty()) {
                throw new RuntimeException("Python脚本未返回任何输出");
            }

            try {
                return new Gson().fromJson(output, JsonObject.class);
            } catch (JsonSyntaxException e) {
                throw new RuntimeException("无效的JSON格式: " + output, e);
            }

        } catch (JSchException | SftpException | InterruptedException e) {
            throw new RuntimeException("SSH operation failed", e);
        } finally {
            // 5. 清理资源
            if (execChannel != null) execChannel.disconnect();
            if (sftpChannel != null) sftpChannel.disconnect();
            if (session != null) session.disconnect();
        }
    }
}
